from langchain_core.prompts import  PromptTemplate,FewShotPromptTemplate
from langchain_openai import ChatOpenAI,OpenAIEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_community.vectorstores import FAISS
from langchain_core.example_selectors import SemanticSimilarityExampleSelector



ZHIPUAI_API_KEY = '794380a4cee054a0f96bb2844b41fd12.X4t70kph1CfmoKfT'
BASE_PATH ='https://open.bigmodel.cn/api/paas/v4/'
parser = StrOutputParser()
 

model = ChatOpenAI(model_name='glm-4',temperature=.7,openai_api_key=ZHIPUAI_API_KEY,base_url=BASE_PATH)

example_prompt = PromptTemplate(input_variables=["input","output"],template="示例输入：{input},输出：{output}")
examples = [
    {"input":"海盗","output":"船"},
    {"input":"飞行员","output":"飞机"},
    {"input":"驾驶员","output":"汽车"},
    {"input":"树","output":"地面"},
    {"input":"鸟","output":"鸟巢"}
]

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
from langchain_community.embeddings import HuggingFaceEmbeddings
# pip install -U langchain-huggingface
# from langchain_huggingface import HuggingFaceEmbeddings
embeddings=HuggingFaceEmbeddings(model_name='distilbert/distilbert-base-uncased')
example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    embeddings,
    FAISS,
    # k=2
)
similar_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    # examples=examples,
    example_prompt=example_prompt,
    prefix="根据下面示例写出输出",
    suffix="你的输入：{what}，输出：",
    input_variables=["what"],
    example_separator="\n"
)
what = "学生"
print(similar_prompt.format(what=what))
print(model.invoke(similar_prompt.format(what=what)).content)
